package cn.iocoder.yudao.framework.server.websocket.security;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.iocoder.yudao.framework.cache.NettyChannelCacheService;
import cn.iocoder.yudao.framework.cache.NettyWsUserChannelCacheService;
import cn.iocoder.yudao.framework.common.enums.UserTypeEnum;
import cn.iocoder.yudao.framework.handlermapping.NettyWsChannelContext;
import cn.iocoder.yudao.framework.handlermapping.WsSessionContext;
import cn.iocoder.yudao.framework.prometheus.ws.WebSocketConnectionMetrics;
import cn.iocoder.yudao.framework.server.websocket.event.WebSocketAuthSuccessEvent;
import cn.iocoder.yudao.framework.server.websocket.WebSocketServerChannelInitializer;
import cn.iocoder.yudao.module.infra.util.TraceIdUtil;
import cn.iocoder.yudao.module.system.api.oauth2.OAuth2TokenApi;
import cn.iocoder.yudao.module.system.api.oauth2.dto.OAuth2AccessTokenCheckRespDTO;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.handler.codec.http.*;
import io.netty.util.Attribute;
import io.netty.util.AttributeKey;
import lombok.extern.slf4j.Slf4j;

import jakarta.annotation.Resource;

import java.util.Objects;

@Slf4j
public class WebTokenAuthorizationHandler extends ChannelInboundHandlerAdapter {
    @Resource
    private OAuth2TokenApi oauth2TokenApi;

    @Resource
    private WebSocketConnectionMetrics webSocketConnectionMetrics;


    private WebSocketUser checkAccessToken(String token) {
        token = SecurityFrameworkUtils.obtainAuthorization(token);
        if (StrUtil.isEmpty(token)) {
            log.warn("token为空, token = {}", token);
            return null;
        }

        final OAuth2AccessTokenCheckRespDTO accessToken = oauth2TokenApi.checkAccessToken(token);
        if (accessToken == null) {
            log.warn("未查询到accessToken, token = {}", token);
            return null;
        }

        final UserTypeEnum userType = UserTypeEnum.valueOf(accessToken.getUserType());
        if (userType != UserTypeEnum.MEMBER) {
            log.warn("不支持的用户类型, userType = {}, token = {}", userType, token);
            return null;
        }

        final WebSocketUser user = new WebSocketUser();
        user.setUserId(accessToken.getUserId())
                .setUserType(accessToken.getUserType())
                .setUserInfo(accessToken.getUserInfo())
                .setTenantId(accessToken.getTenantId())
                .setScopes(accessToken.getScopes());
        return user;
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        final WebSocketUser existedUser = (WebSocketUser) ctx.channel().attr(AttributeKey.valueOf(SecurityFrameworkUtils.WEB_SOCKET_USER_ATTRIBUTE_KEY)).get();
        if (existedUser == null && (msg instanceof FullHttpMessage)) {
            // 握手时进行token鉴权
            try {
                TraceIdUtil.setTraceId(IdUtil.objectId());

                // 记录连接数
                webSocketConnectionMetrics.incrementConnections();

                final FullHttpMessage httpMessage = (FullHttpMessage) msg;
                final HttpHeaders headers = httpMessage.headers();

                // 开启鉴权 认证
                final String token = Objects.requireNonNull(headers.get(SecurityFrameworkUtils.TOKEN_HEADER));
                final WebSocketUser user = checkAccessToken(token);
                this.bindWebSocketSession(user, ctx, headers);
                log.info("Token验证成功, userId = {}, sessionId = {}", user.getUserId(), ctx.channel().id().asLongText());
                ctx.fireChannelRead(msg);
            } catch (Exception e) {
                refuseChannel(ctx);
            } finally {
                TraceIdUtil.removeTraceId();
            }
        } else {
            //other protocols
            // 发送消息时只验证channel的鉴权结果
            if (existedUser == null) {
                refuseChannel(ctx);
                return;
            }

            ctx.fireChannelRead(msg);
        }
    }

    private void bindWebSocketSession(WebSocketUser user, ChannelHandlerContext ctx, HttpHeaders headers) {
        final Attribute<WsSessionContext> attr = ctx.channel().attr(NettyWsChannelContext.WEB_SOCKET_USER_BIND_ATTRIBUTE_KEY);

        final WsSessionContext context = new WsSessionContext();
        context.setWebSocketUser(user);
        context.setSessionId(ctx.channel().id().asLongText());
        attr.set(context);

        ctx.channel().attr(AttributeKey.valueOf(SecurityFrameworkUtils.WEB_SOCKET_USER_ATTRIBUTE_KEY)).set(user);
        TraceIdUtil.setUserId(user.getUserId());

        // 认证通过绑定用户id和当前连接
        final NettyWsUserChannelCacheService userChannelCacheService = SpringUtil.getBean(NettyWsUserChannelCacheService.class);
        userChannelCacheService.bindChannel(user, ctx);
        final NettyChannelCacheService nettyChannelCacheService = SpringUtil.getBean(NettyChannelCacheService.class);
        nettyChannelCacheService.addChannel(ctx, context);

        // 发送认证成功事件
        final WebSocketAuthSuccessEvent authSuccessEvent = new WebSocketAuthSuccessEvent(this, headers);
        authSuccessEvent.setContext(context);
        authSuccessEvent.setWebSocketUser(user);

        final String traceId = TraceIdUtil.getTraceId();
        WebSocketServerChannelInitializer.businessExecutors.submit(() -> {
            try {
                TraceIdUtil.setUserId(user.getUserId());
                TraceIdUtil.setTraceId(traceId);
                SpringUtil.publishEvent(authSuccessEvent);
            } finally {
                TraceIdUtil.removeUserId();
                TraceIdUtil.removeTraceId();
            }
        });
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        super.exceptionCaught(ctx, cause);
    }

    private void refuseChannel(ChannelHandlerContext ctx) {
        ctx.channel().writeAndFlush(new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.UNAUTHORIZED));
        ctx.channel().close();
    }

}
